from backend.api import (
    get_factor_corr,
    get_factor_list,
    get_factor_values,
    get_kline_data_by_pivot,
    get_symbol_list,
)
from backend.draw import draw_kline, draw_one_factor_pic
import pandas as pd
import streamlit as st


if __name__ == "__main__":
    st.title("因子任务")
    with st.form("factor_calculate"):
        main_col, side_col = st.columns([4, 2])
        with main_col:
            timeframe = st.selectbox("Timeframe", ["1m", "5m", "10m"], index=0)
            symbols = st.multiselect("Symbol", get_symbol_list(timeframe))
            leftcol, rightcol = st.columns([1, 1])
            with leftcol:
                start = st.date_input("Start Date")
            with rightcol:
                end = st.date_input("End Date")
            factors = st.multiselect("Factor", get_factor_list(timeframe))
        with side_col:
            agg_method = st.selectbox(
                "Aggregate Method", ["因子聚合", "品种聚合", "表格"], index=0
            )
            is_draw_kline = st.checkbox("Draw Kline")
            is_draw_factor = st.checkbox("Draw Factor")
            is_factor_rank = st.checkbox("Factor Rank")
            sbm = st.form_submit_button("刷新", type="primary")
            sbm2 = st.form_submit_button("清除缓存", type="secondary")
            if sbm2:
                get_factor_corr.clear()
                get_factor_list.clear()
                get_factor_values.clear()
                get_kline_data_by_pivot.clear()
                get_symbol_list.clear()

    if sbm and len(symbols) > 0 and end >= start:
        print("submit and do works")
        if is_draw_kline:
            st.info("to draw kline")
            for symbol in symbols:
                df_kline = get_kline_data_by_pivot(symbol, start, end, timeframe)
                if df_kline is None or len(df_kline) <= 0:
                    st.error(f"没有k线数据: {symbol}")
                    continue
                st.plotly_chart(draw_kline(df_kline, symbol))

        if is_draw_factor and len(factors) > 0:
            factor_values = get_factor_values(symbols, factors, start, end, timeframe)
            # 根据factor_values的值，使用streamlit在前端页面中绘制因子曲线图
            if factor_values is None or len(factor_values) <= 0:
                st.error(f"没有因子数据")
            if len(factor_values) > 0:
                if agg_method == "品种聚合":
                    for factor in factors:
                        st.subheader(f"{factor} Factor")
                        p = draw_one_factor_pic(factor_values, factor=factor)
                        if not p:
                            st.error(f"没有因子数据: {factor}")
                        st.plotly_chart(p)
                elif agg_method == "因子聚合":
                    for symbol in symbols:
                        st.subheader(f"{symbol} Symbol")
                        p = draw_one_factor_pic(factor_values, symbol=symbol)
                        if not p:
                            st.error(f"没有因子数据: {symbol}")
                        st.plotly_chart(p)
                else:
                    st.dataframe(pd.DataFrame(factor_values))
        if is_factor_rank:
            # timeframe, symbols, factors, start, end, delay: int = 5
            factor_rank = get_factor_corr(timeframe, symbols, factors, start, end)
            st.dataframe(pd.DataFrame(factor_rank))
